from fastapi import FastAPI, Request
from sse_starlette.sse import EventSourceResponse
import asyncio
import json
import uuid
from pydantic import BaseModel
from typing import Optional
import uvicorn
import inspect
app = FastAPI()
mcpHub = {}
class McpRequest(BaseModel):
    id: Optional[int] = None
    jsonrpc: str
    method: str
    params: Optional[dict] = None
class MCPServer:
    def __init__(self, name, message_path, tools):
        self.queue = asyncio.Queue()
        self.client_id = str(uuid.uuid4())
        self.message_path = message_path
        self.info = {
            "protocolVersion": "2024-11-05",
            "capabilities": {
                "experimental": {},
                "tools": {
                    "listChanged": False
                }
            },
            "serverInfo": {
                "name": name,
                "version": "1.6.0"
            }
        }
        self.tools = tools
    def list_tool(self):
        result = []
        for tool in self.tools:
            toolInfo = {
                "name": tool.__name__,
                "description": tool.__doc__,
                "inputSchema": {"type": "object","properties":{}},
            }
            for name, param in inspect.signature(tool).parameters.items():
                toolInfo["inputSchema"]["properties"][name] = {
                    "title": name,
                    "type": "string",
                }
            result.append(toolInfo)
        return result
    async def reader(self):
        while True:
            event = await self.queue.get()
            yield event
    @staticmethod
    def response(result, id):
        message = {
            "jsonrpc": "2.0",
            "result": result,
        }
        if id is not None:
            message["id"] = id
        return json.dumps(message)
    async def request(self, req: McpRequest):
        if req.method == "initialize":
            await self.queue.put({"event": "message", "data": self.response(self.info, req.id)})
        elif req.method == "tools/list":
            await self.queue.put({"event": "message", "data": self.response({"tools": self.list_tool()}, req.id)})
        elif req.method == "tools/call":
            for tool in self.tools:
                if tool.__name__ == req.params.get("name"):
                    result = await tool(**req.params["arguments"])
                    await self.queue.put({"event": "message", "data": self.response({"content": result, "isError": False}, req.id)})
                    break
async def test(state=None):
    """
    description
    """
    result = f"hi {state}"
    await asyncio.sleep(1)
    result += "!"
    return result
@app.get("/receive_test")
async def receive_test():
    mcp = MCPServer(name="mcp-test",message_path="/send_test", tools=[test])
    mcpHub[mcp.client_id] = mcp
    await mcp.queue.put({"event": "endpoint", "data": f"{mcp.message_path}?client_id={mcp.client_id}"})
    return EventSourceResponse(mcp.reader())
@app.post("/send_test")
async def send_test(request: Request, payload: McpRequest):
    client_id = request.query_params.get("client_id")
    if client_id not in mcpHub:
        return "no client"
    await mcpHub[client_id].request(payload)
    return "ok"
if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8001)